#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/1/17 9:28
# @Author  : Zangzihan
# @File    : onnx2trt.py
# @Description : 这个把onnx模型转成tensorrt模型
import tensorrt as trt
import os

EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)


def GiB(val):
    return val * 1 << 30


def build_engine(onnx_file_path, engine_file_path, half=False):
    """Takes an ONNX file and creates a TensorRT engine to run inference with"""
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
    with trt.Builder(TRT_LOGGER) as builder, \
            builder.create_network(EXPLICIT_BATCH) as network, \
            builder.create_builder_config() as config, \
            trt.OnnxParser(network, TRT_LOGGER) as parser, \
            trt.Runtime(TRT_LOGGER) as runtime:
        config.max_workspace_size = GiB(6)
        builder.max_batch_size = 12
        half &= builder.platform_has_fast_fp16
        if half:
            config.set_flag(trt.BuilderFlag.FP16)
        # Parse model file
        if not os.path.exists(onnx_file_path):
            print('ONNX file {} not found.'.format(onnx_file_path))
            exit(0)
        print('Loading ONNX file from path {}...'.format(onnx_file_path))
        with open(onnx_file_path, 'rb') as model:
            print('Beginning ONNX file parsing')
            if not parser.parse(model.read()):
                print('ERROR: Failed to parse the ONNX file.')
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                return None
        print('Completed parsing of ONNX file')
        print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
        # profile = builder.create_optimization_profile()
        # profile.set_shape("images", (1,3,640,640), (8,3,640,640), (12,3,640,640))
        # idx = config.add_optimization_profile(profile)
        plan = builder.build_serialized_network(network, config)
        engine = runtime.deserialize_cuda_engine(plan)
        print("Completed creating Engine")
        with open(engine_file_path, "wb") as f:
            f.write(plan)
        return engine


if __name__ == '__main__':
    onnx_path = "logs/20240116-shougang-cicada/shougang-cicada-240117.onnx"
    tensorrt_path = "logs/20240116-shougang-cicada/shougang-cicada-240117.trt"
    build_engine(onnx_path, tensorrt_path)
